package com.ktjiaoyu.travelaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.ktjiaoyu.travelaiagent.agent.enums.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础Agent类，用于管理状态和执行流程。
 *
 * 提供状态转换、记忆管理和基于步骤的执行循环的基础功能。
 * 子类须实现step抽象方法。
 */
@Data
@Slf4j
public abstract class BaseAgent {
    // Agent唯一标识，默认产生一个新对话ID
    protected String chatId = UUID.randomUUID().toString();
    // Agent名称
    protected String name;
    // 系统提示词
    private String systemPrompt;
    // 下一个步骤提示词
    private String nextStepPrompt;
    // 大模型客户端（对应OpenManus中的LLM对象）
    protected ChatClient chatClient;
    // Memory（需要自主维护会话记忆），默认使用内存记忆
    protected ChatMemory chatMemory = new InMemoryChatMemory();
    // Agent状态
    private AgentState state = AgentState.IDLE;
    // 最大执行步骤
    private int maxSteps = 10;
    // 当前执行步骤
    private int currentStep = 0;

    /** 执行一个步骤，要求子类必须实现 */
    public abstract String step() ;
    /** 运行Agent智能体，要求用户输入提示词 */
    public String run(String userPrompt) {
        if (this.state != AgentState.IDLE) {
            // 必须是空闲状态才能运行
            throw new RuntimeException("Cannot run agent from state: " + state);
        }
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty userPrompt");
        }
        // 开始记录消息上下文
        chatMemory.add(chatId, new UserMessage(userPrompt));
        // 保存结果列表
        List<String> results = new ArrayList<>();
        // 更改状态为运行中
        state = AgentState.RUNNING;
        try {
            // 执行循环达到最大步骤或者已经完成为止
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                // 记录当前步骤
                currentStep = i + 1;
                log.info("Executing step {}/{}", currentStep, maxSteps);

                // 执行当前步骤，并记录结果
                String stepResult = step();
                String result = "Step " + currentStep + ": " + stepResult;
                results.add(result);
            }
            // 超出步骤限制则添加终止原因
            if (currentStep >= maxSteps) {
                // 为了简单，暂时不考虑超出最大步骤重置
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + maxSteps + ")");
            }
        }catch (Exception e) {
            // 更改状态为异常
            state = AgentState.ERROR;
            results.add( "执行出现异常：" + e.getMessage());
            log.error("Agent run error: {}", e.getMessage());
        }finally {
            // 清理资源
            cleanResources();
        }
        // 返回整个Agent的执行结果
        return String.join("\n", results);
    }

    /**
     * 增加流式输出的运行方法
     *
     * @param userPrompt 用户提示词
     * @return SseEmitter实例
     */
    public SseEmitter streamRun(String userPrompt) {
        // 创建SseEmitter，设置较长的超时时间
        SseEmitter emitter = new SseEmitter(300000L); // 多次调用AI，设置时间长一点

        // 使用线程异步处理，避免阻塞主线程
        CompletableFuture.runAsync(() -> {
            try {
                if (this.state != AgentState.IDLE) {
                    // 前台提示中文比较好
                    emitter.send("错误：无法从状态运行代理: " + this.state);
                    emitter.complete();
                    return;
                }
                if (StrUtil.isBlank(userPrompt)) {
                    emitter.send("错误：不能使用空提示词运行代理");
                    emitter.complete();
                    return;
                }
                // 记录消息上下文
                chatMemory.add(chatId, new UserMessage(userPrompt));
                // 更改状态为忙碌
                state = AgentState.RUNNING;
                try {
                    for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                        currentStep = i + 1;
                        log.info("Executing step {}/{}", currentStep, maxSteps);

                        String result = "Step " + currentStep + ": " + step();
                        // 发送每一步的结果
                        emitter.send(result);
                    }
                    // 检查是否超出步骤限制
                    if (currentStep >= maxSteps) {
                        state = AgentState.FINISHED;
                        emitter.send("执行结束: 达到最大步骤 (" + maxSteps + ")");
                    }
                    // 完成之前，获取最后一次有效的AI响应并发送
                    if (state == AgentState.FINISHED) {
                        String finalAiResponse = getLastValidAiResponse();
                        // 发送最终响应给前端
                        if (StrUtil.isNotBlank(finalAiResponse)) {
                            emitter.send("\n=== 智能体最终响应 ===");
                            // 发送最终响应
                            emitter.send(finalAiResponse);
                            emitter.send("=== 响应结束 ===");
                        }
                    }
                    // 正常完成
                    emitter.complete();
                } catch (Exception e) {
                    state = AgentState.ERROR;
                    try {
                        emitter.send("执行出现异常: " + e.getMessage());
                        emitter.complete();
                    } catch (Exception ex) {
                        emitter.completeWithError(ex);
                    }
                } finally {
                    // 清理资源
                    this.cleanResources();
                }
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        // 设置超时回调
        emitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanResources();
            log.warn("SSE connection timed out");
        });
        // 设置完成回调
        emitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanResources();
            log.info("SSE connection completed");
        });
        return emitter;
    }

    protected void cleanResources() {
        // 子类可以重写此方法来清理资源
    }

    /**
     * 获取最后一次有效的AI响应
     * 从消息列表中反向查找最后一个助手消息（AssistantMessage）
     *
     * @return 最后一次AI的有效响应文本，如果没有则返回空字符串
     */
    protected String getLastValidAiResponse() {
        // 获取所有消息
        List<Message> messageList = chatMemory.get(this.chatId, Integer.MAX_VALUE);
        // 如果消息为空，则最后的响应也为空
        if (messageList == null || messageList.isEmpty()) {
            return "";
        }
        // 从后往前遍历消息列表，查找最后一个助手消息
        for (int i = messageList.size() - 1; i >= 0; i--) {
            Message message = messageList.get(i);
            // 获取最后一条助手消息
            if (message instanceof AssistantMessage) {
                AssistantMessage assistantMessage = (AssistantMessage) message;
                String content = assistantMessage.getText();
                // 返回非空的助手消息内容
                if (StrUtil.isNotBlank(content)) {
                    log.info("Final response : {}" , content);
                    return content;
                }
            }
        }
        return "智能体未返回最终响应";
    }
}
